package com.hulk.rpc.registry.impl.zk;

import com.google.common.base.Preconditions;
import com.hulk.rpc.config.common.AddressWithWeight;
import com.hulk.rpc.config.common.HostPort;
import com.hulk.rpc.config.server.ProtocolType;
import com.hulk.rpc.registry.ServerMetaInfo;
import com.hulk.rpc.registry.ServerRegistry;
import com.hulk.rpc.common.ForeverRetryPolicy;
import com.hulk.rpc.utils.HexUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCache.StartMode;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.zookeeper.CreateMode;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;

@Slf4j
public class ZooKeeperRegister implements ServerRegistry {

  private CuratorFramework client;
  private ConcurrentMap<String, PathChildrenCache> watcherMap;

  @Override
  public void init(List<HostPort> hostPorts) {
    watcherMap = new ConcurrentHashMap<>();
    String connectString = hostPorts.stream().map(HostPort::toString).collect(Collectors.joining(","));
    RetryPolicy retryPolicy = new ForeverRetryPolicy(1000, 60 * 1000);
    client = CuratorFrameworkFactory.newClient(connectString, 1000 * 10, 1000 * 3, retryPolicy);
    client.start();
  }

  @Override
  public void register(ServerMetaInfo serverMetaInfo) {
    Preconditions.checkNotNull(client);

    String path = StringUtils.join("/hulk/", serverMetaInfo.getGroup(),
            "/", serverMetaInfo.getApp(), "/", serverMetaInfo.getProtocolType().name(),
            "/", HexUtils.toHex(serverMetaInfo.getHostPort().toString().getBytes(StandardCharsets.UTF_8)));
    HostPort serverAddress = serverMetaInfo.getHostPort();
    Integer serverWeight = serverMetaInfo.getWeight();

    byte[] data = new AddressWithWeight(serverAddress, serverWeight).toBytes();

    try {
      if (client.checkExists().forPath(path) != null) {
        if (log.isInfoEnabled()) {
          log.info("删除zk已存在节点: " + path + ", " + serverAddress);
        }

        client.delete().forPath(path);
      }
    } catch (Exception e) {
      if (log.isErrorEnabled()) {
        log.error("zk已存在节点删除失败, " + path + ", " + serverAddress, e);
      }
    }

    try {
      client//
              .create()//
              .creatingParentsIfNeeded()//
              .withMode(CreateMode.EPHEMERAL)//
              .forPath(path, data);

      if (log.isInfoEnabled()) {
        log.info("zk注册成功, " + path + ", " + serverAddress + "@" + serverWeight);
      }
    } catch (Exception e) {
      if (log.isErrorEnabled()) {
        log.error("zk注册失败, " + path + ", " + serverAddress + "@" + serverWeight, e);
      }
    }

    if (!watcherMap.containsKey(path)) {
      addRegisterWatcher(serverMetaInfo.getGroup(), serverMetaInfo.getApp(), serverMetaInfo.getProtocolType(), serverAddress, serverWeight);
    }
  }

  @Override
  public void unregister(ServerMetaInfo serverMetaInfo) {
    String path = StringUtils.join("/hulk/", serverMetaInfo.getGroup(),
            "/", serverMetaInfo.getApp(), "/", serverMetaInfo.getProtocolType().name(),
            "/", HexUtils.toHex(serverMetaInfo.getHostPort().toString().getBytes(StandardCharsets.UTF_8)));
    HostPort serverAddress = serverMetaInfo.getHostPort();

    try {
      PathChildrenCache watcher = watcherMap.remove(path);
      if (watcher != null) {
        watcher.close();
      }
    } catch (Exception e) {
      if (log.isErrorEnabled()) {
        log.error("warcher关闭失败, " + path + ", " + serverAddress, e);
      }
    }

    try {
      if (client.checkExists().forPath(path) != null) {
        client.delete().forPath(path);
      }

      if (log.isInfoEnabled()) {
        log.info("zk注销成功, " + path + ", " + serverAddress);
      }
    } catch (Exception e) {
      if (log.isErrorEnabled()) {
        log.error("zk注销失败, " + path + ", " + serverAddress, e);
      }
    }

  }

  /**
   * 当断掉连接又重新连接上时起作用
   */
  private synchronized void addRegisterWatcher(String group, String app, ProtocolType protocol, HostPort serverAddress,
                                               int serverWeight) {
    final String path = "/hulk/" + group + "/" + app + "/" + protocol + "/"
            + HexUtils.toHex(serverAddress.toString().getBytes(StandardCharsets.UTF_8));

    if (watcherMap.containsKey(path)) {
      return;
    }

    PathChildrenCache watcher = new PathChildrenCache(client, path, false);

    PathChildrenCacheListener pathChildrenCacheListener = new PathChildrenCacheListener() {
      private volatile boolean waitForInitializedEvent = true;

      @Override
      public void childEvent(CuratorFramework client, PathChildrenCacheEvent event) throws Exception {
        switch (event.getType()) {

          case INITIALIZED:
            waitForInitializedEvent = false;
            break;

          case CONNECTION_RECONNECTED:
            if (waitForInitializedEvent) {
              return;
            }

            if (log.isInfoEnabled()) {
              log.info("获得zk连接尝试重新注册, " + path + ", " + serverAddress + "@" + serverWeight);
            }

            ZooKeeperRegister.this.register(new ServerMetaInfo().setWeight(serverWeight)
                    .setProtocolType(protocol).setHostPort(serverAddress).setApp(app).setGroup(group));
            break;

          default:
            break;
        }
      }
    };

    watcher.getListenable().addListener(pathChildrenCacheListener);

    try {
      watcher.start(StartMode.POST_INITIALIZED_EVENT);
      watcherMap.put(path, watcher);
    } catch (Exception e) {
      if (log.isErrorEnabled()) {
        log.error("zk监听失败, " + path, e);
      }
    }
  }

  @Override
  public void close() throws IOException {
    watcherMap.forEach((path, watcher) -> {
      try {
        watcher.close();
      } catch (Exception e) {
        if (log.isErrorEnabled()) {
          log.error("warcher关闭失败, " + path);
        }
      }
    });

    watcherMap = null;

    client.close();
    client = null;
  }
}
